#include "Mem2Reg.hpp"
#include "IRBuilder.hpp"
#include "Value.hpp"
#include "logging.hpp"
#include <memory>
#include "CodeGenUtil.hpp"

/**
 * @brief Mem2Reg Pass的主入口函数
 * 
 * 该函数执行内存到寄存器的提升过程，将栈上的局部变量提升到SSA格式。
 * 主要步骤：
 * 1. 创建并运行支配树分析
 * 2. 对每个非声明函数：
 *    - 清空相关数据结构
 *    - 插入必要的phi指令
 *    - 执行变量重命名
 * 
 * 注意：函数执行后，冗余的局部变量分配指令将由后续的死代码删除Pass处理
 */
void Mem2Reg::run()
{
    // 创建支配树分析 Pass 的实例
    dominators_ = std::make_unique<Dominators>(m_);
    // 建立支配树
    dominators_->run();
    // 以函数为单元遍历实现 Mem2Reg 算法
    for (auto &f : m_->get_functions())
    {
        if (f.is_declaration())
            continue;
        func_ = &f;
        if (func_->get_basic_blocks().size() >= 1)
        {
            // 预处理alloc指令，适配可达定义分析
            loop_alloc_inv_hoist();
            // 对应伪代码中 phi 指令插入的阶段
            generate_phi();
            // 对应伪代码中重命名阶段
            rename(func_->get_entry_block());

        }
        // 后续 DeadCode 将移除冗余的局部变量的分配空间
    }
}

/**
 * @brief 将所有基本块中的 alloca 指令提取出来，统一移动到函数的入口基本块的起始位置
 */
void Mem2Reg::loop_alloc_inv_hoist(){
    std::vector<Instruction*> alloc_vec;
    for(auto &bb : func_->get_basic_blocks()){
        auto b = &bb;
        std::vector<Instruction*> replace_inst;
        for(auto &inst : b->get_instructions()){
            auto instr = &inst;
            if(instr->is_alloca()){
                alloc_vec.push_back(instr);
                replace_inst.push_back(instr);
            }
        }
        for(auto instr : replace_inst){
            instr->get_parent()->remove_instr(instr);
        }
    }
    auto entry_bb = func_->get_entry_block();
    for(auto &instr : alloc_vec){
        entry_bb->add_instr_begin(instr);
    }
}

/**
 * @brief 在必要的位置插入phi指令
 * 
 * 该函数实现了经典的phi节点插入算法：
 * 1. 收集全局活跃变量：
 *    - 扫描所有store指令
 *    - 识别在多个基本块中被赋值的变量
 * 
 * 2. 插入phi指令：
 *    - 对每个全局活跃变量
 *    - 在其定值点的支配边界处插入phi指令
 *    - 使用工作表法处理迭代式的phi插入
 * 
 * phi指令的插入遵循最小化原则，只在必要的位置插入phi节点
 */
void Mem2Reg::generate_phi()
{
    // 步骤一：找到活跃在多个 block 的全局名字集合，以及它们所属的 bb 块
    // 全局名字就是alloca出来的变量，在 SSA 中具有唯一标识符
    std::set<Value *> alloca_var; // 存储活跃于多个基本块中的 alloca 变量
    // 记录每个 alloca 变量在哪些基本块中被 store 使用
    std::map<Value *, std::set<BasicBlock *>> alloca_var_map;
    for (auto &bb : func_->get_basic_blocks())
    {
        auto b = &bb;
        // alloca变量只会被store使用，只需要排除store的addr是一个gep或者全局变量的情况
        for (auto &inst : b->get_instructions())
        {
            auto instr = &inst;
            if (inst.is_store())
            {
                // 对 store 指令：左值是addr，右值是待存值
                auto l_val = dynamic_cast<StoreInst *>(instr)->get_lval();
                // 如果addr是有效的 alloca 指针（非 gep、非全局变量等）
                if (is_valid_ptr(l_val))
                {
                    alloca_var.insert(l_val); // 收集alloca变量
                    alloca_var_map[l_val].insert(b); // 记录其在哪些基本块中使用
                }
            }
        }
    }
    // TODO1:实现步骤二：从支配树获取支配边界信息，并在对应位置插入 phi 指令
    // 算法流程与文档`Mem2Reg介绍`->`Mem2Reg算法流程`中的`1.插入phi指令`一致
    for (auto var : alloca_var) {
        // 工作表法：记录需要插入phi指令的基本块
        std::set<BasicBlock *> work_list;
        
        // 初始化工作表：将变量定值点的支配边界加入工作表
        for (auto bb : alloca_var_map[var]) {
            auto df = dominators_->get_dominance_frontier(bb);
            work_list.insert(df.begin(), df.end());
        }
        
        // 已插入phi指令的基本块集合
        std::set<BasicBlock *> already_inserted;
        
        // 迭代处理工作表中的基本块
        while (!work_list.empty()) {
            // 从工作表中取出一个基本块
            auto bb = *work_list.begin();
            work_list.erase(work_list.begin());
            
            // 如果该基本块未插入过phi指令
            if (already_inserted.find(bb) == already_inserted.end()) {
                // 在基本块开头插入phi指令
                auto phi = PhiInst::create_phi(var->get_type()->get_pointer_element_type(), bb);
                bb->add_instr_begin(phi);
                phi_map[var] = phi;  // 记录变量和对应的phi指令
                already_inserted.insert(bb);
                
                // 如果该基本块也对变量进行了定值，则需要将其支配边界加入工作表
                if (alloca_var_map[var].find(bb) != alloca_var_map[var].end()) {
                    auto df = dominators_->get_dominance_frontier(bb);
                    for (auto frontier_bb : df) {
                        if (already_inserted.find(frontier_bb) == already_inserted.end()) {
                            work_list.insert(frontier_bb);
                        }
                    }
                }
            }
        }
    }
    // TODO1==============end
}

void Mem2Reg::rename(BasicBlock *bb)
{
    std::vector<Instruction *> deleted_instructions;
    // TODO2:依次实现下面6个步骤以完成变量重命名（需阅读文档`Mem2Reg介绍`->`Mem2Reg算法流程`中的`2.变量重命名`）
    // 步骤三：将 phi 指令作为 lval 的最新定值，lval 即是为局部变量 alloca 出的地址空间
    for (auto &instr : bb->get_instructions()) {
        if (instr.is_phi()) {
            auto phi = dynamic_cast<PhiInst *>(&instr);
            // 找到这个phi指令对应的alloca变量
            for (auto &[var, phi_inst] : phi_map) {
                if (phi_inst == phi) {
                    var_stack[var].push_back(phi);
                    break;
                }
            }
        }
    }

    // 步骤四：用 lval 最新的定值替代对应的load指令
    for (auto &instr : bb->get_instructions()) {
        if (instr.is_load()) {
            auto load = dynamic_cast<LoadInst *>(&instr);
            auto lval = load->get_lval();
            if (is_valid_ptr(lval) && !var_stack[lval].empty()) {
                load->replace_all_use_with(var_stack[lval].back());
            }
        }
    }

    // 步骤五：将 store 指令的 rval，也即被存入内存的值，作为 lval 的最新定值
    for (auto &instr : bb->get_instructions()) {
        if (instr.is_store()) {
            auto store = dynamic_cast<StoreInst *>(&instr);
            auto lval = store->get_lval();
            if (is_valid_ptr(lval)) {
                var_stack[lval].push_back(store->get_rval());
                deleted_instructions.push_back(store);
            }
        }
    }

    // 步骤六：为 lval 对应的 phi 指令参数补充完整
    for (auto succ_bb : bb->get_succ_basic_blocks()) {
        for (auto &instr : succ_bb->get_instructions()) {
            if (instr.is_phi()) {
                auto phi = dynamic_cast<PhiInst *>(&instr);
                // 找到这个phi指令对应的alloca变量
                for (auto &[var, phi_inst] : phi_map) {
                    if (phi_inst == phi && !var_stack[var].empty()) {
                        phi->add_phi_pair_operand(var_stack[var].back(), bb);
                        break;
                    }
                }
            }
        }
    }

    // 步骤七：对 bb 在支配树上的所有后继节点，递归执行 rename 操作
    for (auto succ_bb : dominators_->get_dom_tree_succ_blocks(bb)) {
        rename(succ_bb);
    }

    // 步骤八：pop出 lval 的最新定值
    // 先保存当前基本块中的所有变量的栈顶值
    std::map<Value *, size_t> stack_sizes;
    for (auto &[var, stack] : var_stack) {
        stack_sizes[var] = stack.size();
    }
    
    // 对于每个在当前基本块中被定值的变量，将其栈恢复到进入该基本块时的状态
    for (auto &instr : bb->get_instructions()) {
        if (instr.is_phi()) {
            auto phi = dynamic_cast<PhiInst *>(&instr);
            for (auto &[var, phi_inst] : phi_map) {
                if (phi_inst == phi && var_stack[var].size() > stack_sizes[var]) {
                    var_stack[var].pop_back();
                }
            }
        }
    }
    
    for (auto &instr : bb->get_instructions()) {
        if (instr.is_store()) {
            auto store = dynamic_cast<StoreInst *>(&instr);
            auto lval = store->get_lval();
            if (is_valid_ptr(lval) && var_stack[lval].size() > stack_sizes[lval]) {
                var_stack[lval].pop_back();
            }
        }
    }
    // TODO2===============end

    // 步骤九：清除冗余的指令
    for (auto inst : deleted_instructions)
    {
        bb->erase_instr(inst);
    }
}

